"""
FinBERT 推理引擎
"""
import torch
import logging
from typing import Dict, List, Optional, Union

logger = logging.getLogger(__name__)


class FinbertInference:
    """FinBERT推理引擎"""
    
    def __init__(self, model, tokenizer, device: str):
        self.model = model
        self.tokenizer = tokenizer
        self.device = device
        # 获取标签映射
        self.id2label = model.config.id2label if hasattr(model.config, 'id2label') else {}
        self.label2id = model.config.label2id if hasattr(model.config, 'label2id') else {}
    
    def predict(
        self,
        texts: Union[str, List[str]],
        return_all_scores: bool = False,
        **kwargs
    ) -> List:
        """
        情感分析预测
        
        Args:
            texts: 输入文本或文本列表
            return_all_scores: 是否返回所有类别的分数
            **kwargs: 其他参数（如 max_length, truncation, padding 等）
        
        Returns:
            预测结果字典
        """
        try:
            # 保存原始输入类型
            is_single = isinstance(texts, str)
            if is_single:
                texts = [texts]
            
            # 编码输入
            inputs = self.tokenizer(
                texts,
                return_tensors="pt",
                truncation=True,
                padding=True,
                max_length=kwargs.get("max_length", 512)
            )
            
            # 移动到设备
            if self.device != "cpu":
                inputs = {k: v.to(self.device) for k, v in inputs.items()}
            
            # 推理
            with torch.no_grad():
                outputs = self.model(**inputs)
                logits = outputs.logits
            
            # 处理结果
            results = []
            for i, text in enumerate(texts):
                # 获取预测类别
                predicted_class_id = torch.argmax(logits[i], dim=-1).item()
                predicted_label = self.id2label.get(predicted_class_id, f"class_{predicted_class_id}")
                
                # 计算概率
                probabilities = torch.softmax(logits[i], dim=-1).tolist()
                scores = {self.id2label.get(j, f"class_{j}"): float(prob) for j, prob in enumerate(probabilities)}
                
                result = {
                    # "text": text,
                    "label": predicted_label,
                    "score": float(probabilities[predicted_class_id])
                }
                
                if return_all_scores:
                    result["scores"] = scores
                
                results.append(result)
            
            # 根据原始输入类型返回结果
            return results
            
        except Exception as e:
            logger.error(f"预测失败: {str(e)}")
            raise
    
    def classify(
        self,
        text: str,
        return_all_scores: bool = False,
        **kwargs
    ) -> Dict:
        """
        情感分类（单个文本）
        
        Args:
            text: 输入文本
            return_all_scores: 是否返回所有类别的分数
            **kwargs: 其他参数
        
        Returns:
            分类结果字典
        """
        return self.predict(text, return_all_scores=return_all_scores, **kwargs)
    
    def batch_classify(
        self,
        texts: List[str],
        return_all_scores: bool = False,
        **kwargs
    ) -> List:
        """
        批量情感分类
        
        Args:
            texts: 文本列表
            return_all_scores: 是否返回所有类别的分数
            **kwargs: 其他参数
        
        Returns:
            批量分类结果字典
        """
        return self.predict(texts, return_all_scores=return_all_scores, **kwargs)
    
    def get_labels(self) -> List[str]:
        """
        获取模型支持的情感标签列表
        
        Returns:
            标签列表
        """
        return list(self.id2label.values()) if self.id2label else []

